# 实验通用配置

import torch
torch.set_num_threads(1)

from sklearn.model_selection import (
    StratifiedKFold, KFold,
    StratifiedGroupKFold, GroupKFold
)
from src.models.classification.ml_model import RFClassificationModel, XGBClassificationModel, SVMClassificationModel
from src.models.regression.ml_model import LinearRegressionModel, RFRegressionModel, GBRRegressionModel
from src.models.classification.dl_transformer import TransformerClassifier
from src.models.regression.dl_transformer import TransformerRegressor
from src.models.regression.dl_mlp import MLPRegressor
from src.models.regression.dl_lstm import LSTMRegressor
from src.models.regression.dl_rnn import RNNRegressor
from src.models.classification.dl_alexnet import AlexNetClassifier
from src.models.classification.dl_resnet import ResNetClassifier
from src.models.classification.dl_mobilenet import MobileNetClassifier

# 设备
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 数据来源
DATA = {
    "train": "data/imputed/imputed_dataset_all.pkl",
    "test": "data/imputed/imputed_dataset_all_binned.pkl",
    "sensitive": "data/imputed/imputed_iterative_bayes.pkl"
}

# 标签列
LABEL_COL = "cogimp_label"

# 保存目录
DIR_DICT = {
    "experiment": "results/experiments/{experiment}/{model}",
    "sensitive": "results/sensitive/{experiment}/{model}",
    "subgroup": "results/subgroup/{experiment}/{model}"
}


# 数据分割方法
SPLIT = {
    True: {  # 敏感性分析(不是多重插补，因此没有多份数据)
        "classification": {
            "cv": StratifiedKFold,
            "params": {"n_splits": 10, "shuffle": True},
            "needs_groups": False
        },
        "regression": {
            "cv": KFold,
            "params": {"n_splits": 10, "shuffle": True},
            "needs_groups": False
        }
    },
    False: {  # 正常情况（分组）
        "classification": {
            "cv": StratifiedGroupKFold,
            "params": {"n_splits": 10, "shuffle": True},
            "needs_groups": True
        },
        "regression": {
            "cv": GroupKFold,
            "params": {"n_splits": 10},
            "needs_groups": True
        }
    }
}

# 所有模型字典
MODEL_ZOO = {
    "ml": {
        "classification": {
            "rf": RFClassificationModel,
            "xgb": XGBClassificationModel,
            "svm": SVMClassificationModel,
        },
        "regression": {
            "linear": LinearRegressionModel,
            "rf": RFRegressionModel,
            "gbr": GBRRegressionModel,
        },
    },
    "dl": {
        "classification": {
            "transformer": TransformerClassifier,
            "alex": AlexNetClassifier,
            "res": ResNetClassifier,
            "mobile": MobileNetClassifier,
        },
        "regression": {
            "transformer": TransformerRegressor,
            "mlp": MLPRegressor,
            "lstm": LSTMRegressor,
            "rnn": RNNRegressor,
        },
    },
}